package com.gitee.wsl.net.server

import com.gitee.wsl.net.readMedium
import com.gitee.wsl.net.readUMedium
import io.ktor.network.selector.ActorSelectorManager
import io.ktor.network.sockets.Socket
import io.ktor.network.sockets.aSocket
import io.ktor.network.sockets.openReadChannel
import io.ktor.network.sockets.openWriteChannel
import io.ktor.utils.io.ByteReadChannel
import io.ktor.utils.io.ByteWriteChannel
import io.ktor.utils.io.readByte
import io.ktor.utils.io.readInt
import io.ktor.utils.io.writeByte
import io.ktor.utils.io.writeFully
import io.ktor.utils.io.writeInt
import io.ktor.utils.io.writeShort
import kotlinx.coroutines.CoroutineExceptionHandler
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.ExecutorCoroutineDispatcher
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.asCoroutineDispatcher
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.isActive
import kotlinx.coroutines.launch
import kotlinx.coroutines.runBlocking
import timber.log.Timber
import java.util.concurrent.Executors
import kotlin.math.min


class Network(
    private val server: FileServer,
    private val prefetchKeys: IntArray,
    private val revision: Int,
    private val acknowledgeId: Int,
    private val statusId: Int,
) {

    private val exceptionHandler = CoroutineExceptionHandler { context, throwable ->
        Timber.d ( "${throwable.message} $context" )
    }

    private lateinit var dispatcher: ExecutorCoroutineDispatcher
    private var running = false

    /**
     * Start the server and begin creating a new coroutine for every new connection accepted
     * @param threads a fixed number or 0 to dynamically allocate based on need
     */
    fun start(port: Int, threads: Int) = runBlocking {
        val executor = if (threads == 0) Executors.newCachedThreadPool() else Executors.newFixedThreadPool(threads)
        dispatcher = executor.asCoroutineDispatcher()
        val selector = ActorSelectorManager(dispatcher)
        val supervisor = SupervisorJob()
        val scope = CoroutineScope(coroutineContext + supervisor + exceptionHandler)
        with(scope) {
            val server = aSocket(selector).tcp().bind(port = port)
            running = true
            Timber.d( "Listening for requests on port ${port}..." )
            while (running) {
                val socket = server.accept()
                Timber.i ( "New connection accepted $socket" )
                launch(Dispatchers.IO) {
                    connect(socket)
                }
            }
        }
    }

    suspend fun connect(socket: Socket) {
        val read = socket.openReadChannel()
        val write = socket.openWriteChannel(autoFlush = true)
        synchronise(read, write)
        if (acknowledge(read, write)) {
            Timber.d ( "Client synchronisation complete: $socket" )
            readRequests(read, write)
        }
    }

    /**
     * If the client is up-to-date and in the correct state send it the [prefetchKeys] list so it knows what indices are available to request
     */
    suspend fun synchronise(read: ByteReadChannel, write: ByteWriteChannel) {
        val opcode = read.readByte().toInt()
        if (opcode != SYNCHRONISE) {
            Timber.d( "Invalid sync session id: $opcode" )
            write.writeByte(REJECT_SESSION)
            write.flushAndClose()
            return
        }

        val revision = read.readInt()
        if (revision != this.revision) {
            Timber.d( "Invalid game revision: $revision" )
            write.writeByte(GAME_UPDATED)
            write.flushAndClose()
            return
        }

        write.writeByte(0)
        prefetchKeys.forEach { key ->
            write.writeInt(key)
        }
    }

    /**
     * Confirm the client got our message and is ready to start sending file requests
     */
    suspend fun acknowledge(read: ByteReadChannel, write: ByteWriteChannel): Boolean {
        val opcode = read.readByte().toInt()
        if (opcode != ACKNOWLEDGE) {
            Timber.d( "Invalid ack opcode: $opcode" )
            write.writeByte(REJECT_SESSION)
            write.flushAndClose()
            return false
        }

        return verify(read, write, acknowledgeId)
    }

    /**
     * Confirm a session value send by the client is as the server [expected]
     */
    suspend fun verify(read: ByteReadChannel, write: ByteWriteChannel, expected: Int): Boolean {
        val id = read.readMedium()
        if (id != expected) {
            Timber.d( "Invalid session id expected: $expected actual: $id" )
            write.writeByte(BAD_SESSION_ID)
            write.flushAndClose()
            return false
        }
        return true
    }

    suspend fun readRequests(read: ByteReadChannel, write: ByteWriteChannel) = coroutineScope {
        try {
            while (isActive) {
                readRequest(read, write)
            }
        } finally {
            Timber.d( "Client disconnected: $read" )
        }
    }

    /**
     * Verify status updates and pass requests onto the [server] to fulfill
     */
    suspend fun readRequest(read: ByteReadChannel, write: ByteWriteChannel) {
        when (val opcode = read.readByte().toInt()) {
            STATUS_LOGGED_OUT, STATUS_LOGGED_IN -> verify(read, write, statusId)
            PRIORITY_REQUEST, PREFETCH_REQUEST -> server.fulfill(read, write, opcode == PREFETCH_REQUEST)
            else -> {
                Timber.w ( "Unknown request $opcode." )
                write.flushAndClose()
            }
        }
    }

    fun stop() {
        running = false
        dispatcher.close()
    }

    companion object {
        // Opcodes
        const val PREFETCH_REQUEST = 0
        const val PRIORITY_REQUEST = 1
        const val SYNCHRONISE = 15
        const val STATUS_LOGGED_IN = 2
        const val STATUS_LOGGED_OUT = 3
        const val ACKNOWLEDGE = 6

        // Response codes
        private const val GAME_UPDATED = 6.toByte()
        private const val BAD_SESSION_ID = 10.toByte()
        private const val REJECT_SESSION = 11.toByte()
    }
}

interface DataProvider {
    fun data(index: Int, archive: Int): ByteArray?

    companion object {
        /*operator fun invoke(cache: CacheLibrary) = object : DataProvider {
            override fun data(index: Int, archive: Int) =
                if (index == 255)
                    cache.index255?.readArchiveSector(archive)?.data
                else
                    cache.index(index).readArchiveSector(archive)?.data
        }*/
    }
}

class FileServer(
    private val provider: DataProvider,
    private val versionTable: ByteArray
) {
    /**
     * Fulfills a request by sending the requested files data to the requester
     */
    suspend fun fulfill(read: ByteReadChannel, write: ByteWriteChannel, prefetch: Boolean) {
        val value = read.readUMedium()
        val index = value shr 16
        val archive = value and 0xffff
        val data = data(index, archive) ?: return Timber.d( "Unable to fulfill request $index $archive $prefetch." )
        serve(write, index, archive, data, prefetch)
    }

    /**
     * @return data for an [index]'s [archive] file or [versionTable] when index and archive are both 255
     */
    fun data(index: Int, archive: Int): ByteArray? {
        if (index == 255 && archive == 255) {
            return versionTable
        }
        return provider.data(index, archive)
    }

    /**
     * Writes response header followed by the contents of [data] to [write]
     */
    suspend fun serve(write: ByteWriteChannel, index: Int, archive: Int, data: ByteArray, prefetch: Boolean) {
        val compression = data[0].toInt()
        val size = getInt(data[1], data[2], data[3], data[4]) + if (compression != 0) 8 else 4
        Timber.d( "Serving file $index $archive - $size." )
        write.writeByte(index.toByte())
        write.writeShort(archive.toShort())
        write.writeByte((if (prefetch) compression or 0x80 else compression).toByte())
        serve(write, HEADER, data, OFFSET, size, SPLIT)
    }

    /**
     * Writes [source] [offset] [size] to [write] and starting at [headerSize] inserting a [SEPARATOR] every [split] bytes
     */
    suspend fun serve(write: ByteWriteChannel, headerSize: Int, source: ByteArray, offset: Int, size: Int, split: Int) {
        var length = min(size, split - headerSize)
        write.writeFully(source, offset, length)
        var written = length
        while (written < size) {
            write.writeByte(SEPARATOR)

            length = if (size - written < split) size - written else split - 1
            write.writeFully(source, written + offset, length)
            written += length
        }
    }

    companion object {

        private fun getInt(b1: Byte, b2: Byte, b3: Byte, b4: Byte) = b1.toInt() shl 24 or (b2.toInt() and 0xff shl 16) or (b3.toInt() and 0xff shl 8) or (b4.toInt() and 0xff)

        private const val SEPARATOR = 255.toByte()
        private const val HEADER = 4
        private const val SPLIT = 512
        private const val OFFSET = 1
    }
}